import asyncio
import time
from typing import Optional

import av
import fractions
from av import AudioFrame

from aiortc_custom import AudioStreamTrack



class FileAudioStreamTrack(AudioStreamTrack):

	def __init__(self, file_path: str, target_sample_rate: int = 48000, target_channels: int = 2):
		super().__init__()
		self._container: Optional[av.container.InputContainer] = None
		self._stream: Optional[av.audio.stream.AudioStream] = None
		self._resampler: Optional[av.audio.resampler.AudioResampler] = None
		self._file_path = file_path
		self._target_sample_rate = target_sample_rate
		self._target_channels = target_channels
		self._timestamp = 0
		self._init_av()

	def _init_av(self) -> None:
		self._container = av.open(self._file_path)
		self._stream = next((s for s in self._container.streams if s.type == "audio"), None)
		if self._stream is None:
			raise RuntimeError("所选视频文件不包含音频流")
		self._stream.thread_type = "AUTO"
		self._resampler = av.audio.resampler.AudioResampler(
			format="s16",
			layout="mono" if self._target_channels == 1 else "stereo",
			rate=self._target_sample_rate,
		)

	async def recv(self) -> AudioFrame:
		if self.readyState != "live":
			raise Exception("track ended")
		assert self._container is not None
		assert self._stream is not None
		assert self._resampler is not None

		# Decode packets until we get an audio frame
		ptime = 0.02
		samples_per_packet = int(self._target_sample_rate * ptime)
		while True:
			for packet in self._container.demux(self._stream):
				for frame in packet.decode():
					if isinstance(frame, AudioFrame):
						for out in self._resampler.resample(frame):
							# ensure timing info
							out.sample_rate = self._target_sample_rate
							out.time_base = fractions.Fraction(1, self._target_sample_rate)
							out.pts = self._timestamp
							self._timestamp += out.samples if hasattr(out, "samples") else samples_per_packet
							# pace to realtime roughly by ptime
							await asyncio.sleep(ptime)
							return out
			# If demux ended, seek to start for looping
			try:
				self._container.seek(0)
			except Exception:
				await asyncio.sleep(0.02)


class MicrophoneAudioStreamTrack(AudioStreamTrack):

	def __init__(self, device: Optional[int] = None, sample_rate: int = 48000, channels: int = 1):
		super().__init__()
		self._device = device
		self._sample_rate = sample_rate
		self._channels = channels
		self._chunk_duration = 0.02  # 20ms
		self._frames_per_chunk = int(self._sample_rate * self._chunk_duration)
		self._sd = None
		self._stream = None
		self._queue = asyncio.Queue(maxsize=10)
		self._timestamp = 0
		self._proc = None
		self._reader_thread = None
		self._init_input()

	def _init_input(self) -> None:
		# 抑制ALSA错误信息
		import os
		os.environ['ALSA_DEBUG'] = '0'
		
		# 方法1: 优先尝试 sounddevice (更稳定，错误更少)
		try:
			import sounddevice as sd  # type: ignore
			self._sd = sd
			
			# 检查可用设备
			try:
				devices = sd.query_devices()
				default_input = sd.query_devices(kind='input')
				print(f"使用默认输入设备: {default_input['name']}")
			except Exception:
				pass
			
			self._stream = sd.InputStream(
				device=self._device,
				samplerate=self._sample_rate,
				channels=self._channels,
				dtype="int16",
				blocksize=self._frames_per_chunk,
				callback=self._on_audio,
			)
			self._stream.start()
			print(f"sounddevice麦克风启动成功: {self._sample_rate}Hz, {self._channels}通道")
			return
		except Exception as e2:
			print(f"sounddevice麦克风启动失败: {e2}")
			pass
		
		# 方法2: 尝试 pyaudio (跨平台，更稳定)
		try:
			import pyaudio
			self._pyaudio = pyaudio
			self._p = pyaudio.PyAudio()
			
			# 获取默认输入设备信息
			try:
				default_input = self._p.get_default_input_device_info()
				print(f"使用默认输入设备: {default_input['name']}")
			except Exception:
				pass
			
			self._stream = self._p.open(
				format=pyaudio.paInt16,
				channels=self._channels,
				rate=self._sample_rate,
				input=True,
				frames_per_buffer=self._frames_per_chunk,
				stream_callback=self._on_audio_pyaudio,
			)
			self._stream.start_stream()
			self._sd = None
			print(f"pyaudio麦克风启动成功: {self._sample_rate}Hz, {self._channels}通道")
			return
		except Exception as e1:
			print(f"pyaudio麦克风启动失败: {e1}")
			pass
		
		# 方法3: Linux 下使用 ffmpeg
		import sys
		import subprocess
		import threading
		try:
			if sys.platform == "linux":
				# 尝试多种音频输入方式，与 DeviceDiagnostics.py 保持一致
				# 添加ALSA错误抑制参数
				cmd_options = [
					["ffmpeg", "-loglevel", "quiet", "-f", "alsa", "-i", "default", "-ac", str(self._channels), "-ar", str(self._sample_rate), "-f", "s16le", "-"],
					["ffmpeg", "-loglevel", "quiet", "-f", "pulse", "-i", "default", "-ac", str(self._channels), "-ar", str(self._sample_rate), "-f", "s16le", "-"],
					["ffmpeg", "-loglevel", "quiet", "-f", "alsa", "-i", "hw:0,0", "-ac", str(self._channels), "-ar", str(self._sample_rate), "-f", "s16le", "-"],
				]
				
				cmd = None
				for cmd_try in cmd_options:
					try:
						# 测试命令是否可用
						test_proc = subprocess.run(cmd_try[:-1] + ["-t", "0.1", "-f", "null", "-"], 
												stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, timeout=2)
						if test_proc.returncode == 0:
							cmd = cmd_try
							break
					except (subprocess.TimeoutExpired, FileNotFoundError):
						continue
				
				if cmd is None:
					raise RuntimeError("无法找到可用的音频输入设备")
				self._proc = subprocess.Popen(
					cmd,
					stdin=subprocess.DEVNULL,
					stdout=subprocess.PIPE,
					stderr=subprocess.DEVNULL,
					bufsize=0,
				)
				bytes_per_chunk = self._frames_per_chunk * self._channels * 2
				def _reader():
					try:
						assert self._proc is not None and self._proc.stdout is not None
						while True:
							data = self._proc.stdout.read(bytes_per_chunk)
							if not data:
								break
							if not self._queue.full():
								self._queue.put_nowait(data)
					except Exception:
						pass
				self._reader_thread = threading.Thread(target=_reader, daemon=True)
				self._reader_thread.start()
				self._sd = None
				self._stream = True  # 标记为可用
				return
		except Exception as e3:
			pass
		
		# 所有方法都失败，使用静音
		print(f"初始化麦克风失败，将使用静音音频轨道")
		self._sd = None
		self._stream = None

	def _on_audio(self, indata, frames, time_info, status):
		try:
			if not self._queue.full():
				self._queue.put_nowait(bytes(indata))
		except Exception:
			pass
	
	def _on_audio_pyaudio(self, in_data, frame_count, time_info, status):
		try:
			if not self._queue.full():
				self._queue.put_nowait(in_data)
		except Exception:
			pass
		return (None, self._pyaudio.paContinue)

	async def recv(self) -> AudioFrame:
		if self.readyState != "live":
			raise Exception("track ended")
		# If microphone not available, generate silence
		if self._stream is None:
			await asyncio.sleep(self._chunk_duration)
			frame = AudioFrame(format="s16", layout="mono" if self._channels == 1 else "stereo", samples=self._frames_per_chunk)
			for p in frame.planes:
				p.update(bytes(p.buffer_size))
			frame.pts = self._timestamp
			self._timestamp += self._frames_per_chunk
			frame.sample_rate = self._sample_rate
			frame.time_base = fractions.Fraction(1, self._sample_rate)
			return frame

		# Get recorded bytes
		try:
			data = await asyncio.wait_for(self._queue.get(), timeout=1.0)
		except asyncio.TimeoutError:
			data = bytes(self._frames_per_chunk * self._channels * 2)

		# Validate data length
		expected_length = self._frames_per_chunk * self._channels * 2
		if len(data) != expected_length:
			# Pad or truncate data to expected length
			if len(data) < expected_length:
				data = data + bytes(expected_length - len(data))
			else:
				data = data[:expected_length]

		frame = AudioFrame(format="s16", layout="mono" if self._channels == 1 else "stereo", samples=self._frames_per_chunk)
		for p in frame.planes:
			p.update(data[: p.buffer_size])
		frame.pts = self._timestamp
		self._timestamp += self._frames_per_chunk
		frame.sample_rate = self._sample_rate
		frame.time_base = fractions.Fraction(1, self._sample_rate)
		
		return frame

	def __del__(self):
		try:
			# 清理 pyaudio
			if hasattr(self, '_p') and self._p is not None:
				if hasattr(self, '_stream') and self._stream is not None:
					self._stream.stop_stream()
					self._stream.close()
				self._p.terminate()
			# 清理 sounddevice
			elif self._stream is not None and hasattr(self._stream, 'stop'):
				self._stream.stop()
				self._stream.close()
			# 清理 ffmpeg 进程
			if self._proc is not None:
				self._proc.terminate()
		except Exception:
			pass


